package com.gitee.zycra.file.monitor;

import com.jcraft.jsch.SftpProgressMonitor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.CountDownLatch;

/**
 * Speed limit for file transfer.
 *
 * @author zycra
 * @since 1.0.0
 */
public class SpeedLimitMonitor implements SftpProgressMonitor {
    private static final Logger LOGGER = LoggerFactory.getLogger(SpeedLimitMonitor.class);

    /**
     * The unit time for speed limit.
     *
     * @since 1.0.0
     */
    private static final Integer SLEEP_TIME = 1000;

    /**
     * The transferred buffer size.
     *
     * @since 1.0.0
     */
    private long transferred;

    /**
     * The total file size.
     *
     * @since 1.0.0
     */
    private long fileSize;

    /**
     * The initional file size.
     *
     * @since 1.0.0
     */
    private long initFileSize;

    /**
     * The start transfer time.
     *
     * @since 1.0.0
     */
    private long startTime;

    /**
     * The speed limit for transfer(kb/s).
     * <p>0 or negative means no limit.
     *
     * @since 1.0.0
     */
    private final long speedLimit;

    /**
     * Call back label for transfer end.
     *
     * @since 1.0.0
     */
    private final CountDownLatch latch;

    public SpeedLimitMonitor(long speedLimit, CountDownLatch latch) {
        this.speedLimit = speedLimit;
        this.latch = latch;
    }

    /**
     * Initialize method.
     *
     * @param type           0 is upload, 1 is download.
     * @param sourceFilePath the path of the source file.
     * @param targetFilePath the path of the target file.
     * @param fileSize       the file size.
     * @since 1.0.0
     */
    @Override
    public void init(int type, String sourceFilePath, String targetFilePath, long fileSize) {
        LOGGER.info("transfer begin type={}, sourceFilePath={}, targetFilePath={}, fileSize={}",
                type == 0 ? "upload" : "download", sourceFilePath, targetFilePath, fileSize);
        this.startTime = System.currentTimeMillis();
        this.fileSize = fileSize;
        this.transferred = 0L;
    }

    /**
     * Called when each block transfer has completed.
     *
     * @param size the size of current transferred block.
     * @return whether continue to transfer.
     * @since 1.0.0
     */
    @Override
    public boolean count(long size) {
        long now = System.currentTimeMillis();
        long timeElapse = now - startTime;
        if (transferred >= fileSize) {
            return false;
        }
        initFileSize = Math.max(initFileSize, size);
        transferred += size;
        if (timeElapse > 0 && speedLimit > 0 && (transferred - initFileSize) / timeElapse > speedLimit) {
            try {
                Thread.sleep(SLEEP_TIME);
            } catch (InterruptedException e) {
                LOGGER.info("transfer thread sleep error", e);
                Thread.currentThread().interrupt();
            }
        }
        return true;
    }

    /**
     * Called when the transfer has completed.
     *
     * @since 1.0.0
     */
    @Override
    public void end() {
        if (latch != null) {
            latch.countDown();
        }
    }
}
